package middleware

import (
	"bytes"
	"fmt"
	"fiberx/internal/global"
	"io"
	"strconv"
	"time"

	"github.com/gin-gonic/gin"
)

// responseSizeWriter 自定义响应写入器，用于记录响应体的大小
type responseSizeWriter struct {
	gin.ResponseWriter
	size int
}

func (w *responseSizeWriter) Write(b []byte) (int, error) {
	size, err := w.ResponseWriter.Write(b)
	w.size += size
	return size, err
}

// bytesToKB 将字节转换为 kb 并保留两位小数
func bytesToKB(bytes int64) float64 {
	return float64(bytes) / 1024
}

// TrafficStatsMiddleware 中间件，用于统计用户的流量使用情况并以 kb 为单位缓存到 Redis
func TrafficStats() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 这里以 IP 地址作为用户标识，实际应用中可根据需求修改
		userID := c.ClientIP()

		// 统计请求体的大小
		requestSize := int64(0)
		if c.Request.Body != nil {
			body, err := io.ReadAll(c.Request.Body)
			if err == nil {
				requestSize = int64(len(body))
				// 恢复请求体，以便后续处理
				c.Request.Body = io.NopCloser(bytes.NewBuffer(body))
			}
		}

		// 记录响应体的大小
		responseWriter := c.Writer
		responseSizeRecorder := &responseSizeWriter{
			ResponseWriter: responseWriter,
			size:           0,
		}
		c.Writer = responseSizeRecorder

		// 继续处理请求
		c.Next()

		// 计算本次请求的总流量
		totalTraffic := requestSize + int64(responseSizeRecorder.size)

		// 将字节转换为 kb
		totalTrafficKB := bytesToKB(totalTraffic)

		// 将以 kb 为单位的流量数据缓存到 Redis 的有序集合中
		key := fmt.Sprintf("traffic:%s", userID)
		err := global.R.Set(c, key, strconv.FormatFloat(totalTrafficKB, 'f', 2, 64), 30*time.Second).Err()
		if err != nil {
			fmt.Printf("Failed to cache traffic data to Redis: %v\n", err)
		}
		fmt.Printf("Traffic data cached to Redis: %s, %f kb", key, totalTrafficKB)
	}
}
